package org.third.integration.ldap.ldaps;

import java.io.IOException;
import java.net.InetAddress;
import java.net.Socket;
import java.net.SocketAddress;
import java.net.UnknownHostException;
import java.security.KeyManagementException;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.ArrayList;

import javax.net.SocketFactory;
import javax.net.ssl.HttpsURLConnection;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLEngine;
import javax.net.ssl.SSLPeerUnverifiedException;
import javax.net.ssl.SSLSession;
import javax.net.ssl.SSLSocket;
import javax.net.ssl.SSLSocketFactory;
import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactory;
import javax.net.ssl.X509ExtendedTrustManager;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.third.common.utils.StringUtils;

import sun.security.util.HostnameChecker;

/**
 * Implements hostname verification for LDAPS which does not have it Java by default (unlike LDAP with StartTLS).
 */
public class LdapsSocketFactory extends SocketFactory {
    private static final Logger logger = LoggerFactory.getLogger(LdapsSocketFactory.class);
    private static ThreadLocal<ArrayList<String>> threadLocalHostnames = new ThreadLocal<ArrayList<String>>();
    private static ThreadLocal<Boolean> threadLocalValidateCertificate = new ThreadLocal<Boolean>();
    
    
    private static ThreadLocal<String> threadLocalSslPublicKey = new ThreadLocal<String>();
    private static KeyStore trustStore;

    private final SSLSocketFactory defaultSSLSocketFactory;

    public LdapsSocketFactory(SSLSocketFactory defaultSSLSocketFactory) {
        this.defaultSSLSocketFactory=defaultSSLSocketFactory;
    }

    public static void setThreadLocalHostnames(ArrayList<String> threadLocalHostnames) {
        LdapsSocketFactory.threadLocalHostnames.set(threadLocalHostnames);
    }

   

	public static void setThreadLocalSslPublicKey(String threadLocalSslPublicKey) {
		LdapsSocketFactory.threadLocalSslPublicKey.set(threadLocalSslPublicKey);
	}

	public static void setThreadLocalValidateCertificate(Boolean validateCertificate) {
        LdapsSocketFactory.threadLocalValidateCertificate.set(validateCertificate);
    }

    public static KeyStore getTrustStore() {
        return trustStore;
    }

    public static void setTrustStore(KeyStore trustStore) {
        LdapsSocketFactory.trustStore = trustStore;
    }

    @Override
    public Socket createSocket() throws IOException {
        // for ldap connect timeout where createSocket() returns unconnected instance which has method connect(java.net.SocketAddress, int timeout) and accepts java.net.InetSocketAddress as first arg
        // validation still needs to occur
        final SSLSocket socket = (SSLSocket) defaultSSLSocketFactory.createSocket();
        return new SSLSocketAbstractDelegate(socket) {
            @Override
            public void connect(SocketAddress endpoint) throws IOException {
                super.connect(endpoint);
                validate(socket);
            }

            @Override
            public void connect(SocketAddress endpoint, int timeout) throws IOException {
                super.connect(endpoint, timeout);
                validate(socket);
            }
        };
    }

    @Override
    public Socket createSocket(String host, int port) throws IOException, UnknownHostException {
        Socket socket = defaultSSLSocketFactory.createSocket(host, port);
        validate((SSLSocket) socket);
        return socket;
    }

    @Override
    public Socket createSocket(String host, int port, InetAddress localHost, int localPort) throws IOException, UnknownHostException {
        Socket socket = defaultSSLSocketFactory.createSocket(host, port, localHost, localPort);
        validate((SSLSocket)socket);
        return socket;
    }

    @Override
    public Socket createSocket(InetAddress host, int port) throws IOException {
        Socket socket = defaultSSLSocketFactory.createSocket(host, port);
        validate((SSLSocket) socket);
        return socket;
    }

    @Override
    public Socket createSocket(InetAddress address, int port, InetAddress localAddress, int localPort) throws IOException {
        Socket socket = defaultSSLSocketFactory.createSocket(address, port, localAddress, localPort);
        validate((SSLSocket) socket);
        return socket;
    }
    
    private void validate(SSLSocket socket) throws IOException {
        ArrayList<String> hostnames = threadLocalHostnames.get();
        if (hostnames == null) return;
        SSLSession session = socket.getSession();
        boolean ok = false;
        for (String hostname : hostnames) {
            HostnameChecker checker = HostnameChecker.getInstance(HostnameChecker.TYPE_TLS);
            try {
                checker.match(hostname, (X509Certificate) session.getPeerCertificates()[0]);
                ok = true;
                break;
            } catch (CertificateException e) {
            }
        }
        threadLocalHostnames.set(null);
        if (!ok) {
            socket.close();
            throw new SSLPeerUnverifiedException(
                                    "Expected hostnames of the server '" + hostnames +
                                    "' do not match the hostname in the " +
                                    "server's certificate.");
        }
    }

    public static SocketFactory getDefault() {
        Boolean validateCertificate = threadLocalValidateCertificate.get();
        String sslPublicKey = threadLocalSslPublicKey.get();
        if (validateCertificate != null && !validateCertificate) {
            //return new LdapsSocketFactory(HttpsURLConnection.getDefaultSSLSocketFactory());
            TrustManager[] trustAllCerts = new TrustManager[]{new X509ExtendedTrustManager() {
                @Override
                public void checkClientTrusted(X509Certificate[] x509Certificates, String s) throws CertificateException {
                }
                @Override
                public void checkServerTrusted(X509Certificate[] x509Certificates, String s) throws CertificateException {
                }
                @Override
                public X509Certificate[] getAcceptedIssuers() {
                    return null;
                }
                @Override
                public void checkClientTrusted(X509Certificate[] x509Certificates, String s, Socket socket) throws CertificateException {
                }
                @Override
                public void checkServerTrusted(X509Certificate[] x509Certificates, String s, Socket socket) throws CertificateException {
                }
                @Override
                public void checkClientTrusted(X509Certificate[] x509Certificates, String s, SSLEngine sslEngine) throws CertificateException {
                }
                @Override
                public void checkServerTrusted(X509Certificate[] x509Certificates, String s, SSLEngine sslEngine) throws CertificateException {
                }
            }};
            SSLContext sc = null;
            // Get a new SSL context
            try {
                sc = SSLContext.getInstance("TLSv1.2");
                sc.init(null, trustAllCerts, new java.security.SecureRandom());
            } catch (NoSuchAlgorithmException e) {
                logger.error("TLSv1.2 Algorithm not found for SSLContext", e);
            } catch (KeyManagementException e) {
                logger.error("Error Key Management", e);
            }
            return new LdapsSocketFactory(sc.getSocketFactory());
        }
		if (validateCertificate && StringUtils.hasText(sslPublicKey)) {
			try {
				SSLContext ctx = SSLContext.getInstance("TLSv1.2");
				ctx.init(null, new TrustManager[] { new LdapTrustManager(sslPublicKey) },
						new java.security.SecureRandom());
				return new LdapsSocketFactory(ctx.getSocketFactory());

			} catch (NoSuchAlgorithmException e) {
				logger.error("TLSv1.2 Algorithm not found for SSLContext", e);
			} catch (KeyManagementException e) {
				logger.error("Error Key Management", e);
			} catch (KeyStoreException e) {
				logger.error("Error Key Store", e);
			} catch (CertificateException e) {
				logger.error("Error CertificateException", e);
			} catch (IOException e) {
				logger.error("Error IO", e);
			}
		}
        if (trustStore!=null) {
            try {
                TrustManagerFactory tmf = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
                tmf.init(trustStore);
                TrustManager[] trustManagers = tmf.getTrustManagers();
                SSLContext sc = SSLContext.getInstance("TLSv1.2");
                sc.init(null, trustManagers, new java.security.SecureRandom());
                return new LdapsSocketFactory(sc.getSocketFactory());
            } catch (NoSuchAlgorithmException e) {
                logger.error("Error to customize SSL context", e);
            } catch (KeyManagementException e) {
                logger.error("Error Key Management", e);
            } catch (KeyStoreException e) {
                logger.error("Error Key Management", e);
            }
        }
        return new LdapsSocketFactory(HttpsURLConnection.getDefaultSSLSocketFactory());
    }
}
